import sqlite3
import os

# 数据库文件名
DB_NAME = "step5.db"


def create_database():
    """
    创建并初始化SQLite数据库和表。
    如果数据库文件已存在，则不会重复创建。
    """
    if os.path.exists(DB_NAME):
        print(f"数据库 '{DB_NAME}' 已存在，跳过创建。")
        return

    try:
        print(f"正在创建数据库: {DB_NAME}")
        # 连接到数据库（如果不存在，则会创建文件）
        conn = sqlite3.connect(DB_NAME)
        cursor = conn.cursor()

        # 定义建表SQL语句
        # 使用 TEXT 作为主键的一部分来提高查询效率
        # prediction_probability 存储为 REAL 类型
        create_table_query = """
                             CREATE TABLE sentence_predictions \
                             ( \
                                 report_id              TEXT NOT NULL, \
                                 company_name           TEXT, \
                                 year                   INTEGER, \
                                 sentence_text          TEXT NOT NULL, \
                                 predicted_class        INTEGER, \
                                 prediction_probability REAL, \
                                 PRIMARY KEY (report_id, sentence_text)
                             ); \
                             """

        # 执行SQL
        cursor.execute(create_table_query)

        # 创建索引以加速查询
        print("正在创建索引...")
        cursor.execute("CREATE INDEX idx_company_year ON sentence_predictions (company_name, year);")
        cursor.execute("CREATE INDEX idx_predicted_class ON sentence_predictions (predicted_class);")

        # 提交事务并关闭连接
        conn.commit()
        conn.close()
        print("数据库和表'sentence_predictions'创建成功。")

    except sqlite3.Error as e:
        print(f"数据库操作失败: {e}")


if __name__ == "__main__":
    create_database()

